#include <argp.h>
#include <bpf/libbpf.h>
#include <sys/resource.h>
#include <procstat.h>
#include <stdio.h>
#include <signal.h>
#include <procstat.skel.h>
#include <time.h>
#include <sys/resource.h>
#include <unistd.h>
#include <cmdparser.h>

static cmdp_action_t head_process(cmdp_process_param_st* paras);

struct {
	int watch_pid;
} arg_procstat;

cmdp_command_st procstat_cmd = {
	.name = "procstat",
	.desc = "show memory usage of a process",
	.doc = "procstat [<args>]\n"
		   "\n",
	.options =
		(cmdp_option_st[]){
			{'p', "pid", "Show memory usage of a process.", CMDP_TYPE_INT4,
			 &arg_procstat.watch_pid},
			{0},
		},
	.fn_process = head_process,
};

typedef struct {
	struct ring_buffer* rb;
	struct procstat_bpf* skel;
} procstat_obj;

static procstat_obj procstat = {0};
static bool flag = false;

static cmdp_action_t head_process(cmdp_process_param_st* paras) {
	if (paras->opts) {
		flag = true;
		return CMDP_ACT_OK;
	}

	return CMDP_ACT_SHOW_HELP;
}

static int libbpf_print_fn(enum libbpf_print_level level, const char* format, va_list args) {
	return vfprintf(stderr, format, args);
}

static volatile bool exiting = false;

static void sig_handler(int sig) {
	exiting = true;
}

static int handle_event(void* ctx, void* data, size_t data_sz) {
	const struct procstat_event* e = data;
	struct tm* tm;
	char ts[32];
	time_t t;

	time(&t);
	tm = localtime(&t);
	strftime(ts, sizeof(ts), "%H:%M:%S", tm);

	printf("%-8s %-8d %-8ld %-8lld %-8lld %-8lld %-8d\n", ts, e->pid, e->size, e->rssanon,
		   e->rssfile, e->rssshmem, e->ppid);

	return 0;
}

void procstat_run(void) {
	struct ring_buffer* rb = procstat.rb;
	int err;

	printf("%-8s %-8s %-8s %-8s %-8s %-8s\n", "TIME", "PID", "SIZE", "RSSANON", "RSSFILE",
		   "RSSSHMEM");

	while (!exiting) {
		err = ring_buffer__poll(rb, 100 /* timeout, ms */);
		/* Ctrl-C will cause -EINTR */
		if (err == -EINTR) {
			err = 0;
			break;
		}
		if (err < 0) {
			printf("Error polling perf buffer: %d\n", err);
			break;
		}
	}
}

int procstat_release(void) {
	struct ring_buffer* rb = procstat.rb;
	struct procstat_bpf* skel = procstat.skel;

	if (rb) {
		ring_buffer__free(rb);
	}
	if (skel) {
		procstat_bpf__destroy(skel);
	}
	return 0;
}

int procstat_open(int argc, char** argv) {
	struct ring_buffer* rb;
	struct procstat_bpf* skel;
	int own_pid;
	int err;

	if (cmdp_run(argc, argv, &procstat_cmd, NULL) || !flag)
		return -1;

	/* Set up libbpf errors and debug info callback */
	libbpf_set_print(libbpf_print_fn);

	/* Cleaner handling of Ctrl-C */
	signal(SIGINT, sig_handler);
	signal(SIGTERM, sig_handler);

	/* Load and verify BPF application */
	skel = procstat_bpf__open();
	if (!skel) {
		fprintf(stderr, "Failed to open and load BPF skeleton\n");
		return -1;
	}
	own_pid = getpid();
	skel->bss->user_pid = own_pid;
	skel->bss->watch_pid = arg_procstat.watch_pid;
	procstat.skel = skel;
	/* Load & verify BPF programs */
	err = procstat_bpf__load(skel);
	if (err) {
		fprintf(stderr, "Failed to load and verify BPF skeleton\n");
		goto cleanup;
	}

	/* Attach tracepoints */
	err = procstat_bpf__attach(skel);
	if (err) {
		fprintf(stderr, "Failed to attach BPF skeleton\n");
		goto cleanup;
	}

	/* Set up ring buffer polling */
	rb = ring_buffer__new(bpf_map__fd(skel->maps.rb), handle_event, NULL, NULL);
	if (!rb) {
		err = -1;
		fprintf(stderr, "Failed to create ring buffer\n");
		goto cleanup;
	}
	procstat.rb = rb;
	return 0;

cleanup:
	procstat_release();
	return -1;
}
